
import math
import torch
from torch import nn
from modeling.attention.cbam import CBAM, CBAM_fby, CBAM_fby_cat
from modeling.attention.SE import SEAttention as SE
from modeling.attention.SK import SKAttention as SK
from modeling.attention.CA import CoordAtt as CA
from modeling.attention.gc import GlobalContextBlock, GlobalContextBlock_fby
from modeling.backbones.resnet import Bottleneck,BasicBlock_group,BasicBlock

# class ResNet_CBAM(nn.Module):
#     def __init__(self, last_stride=2, block=Bottleneck, layers=[3, 4, 6, 3]):
#         self.inplanes = 64
#         super().__init__()
#         self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3,
#                                bias=False)
#         self.bn1 = nn.BatchNorm2d(64)
#         # self.relu = nn.ReLU(inplace=True)   # add missed relu
#         self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
#         self.layer1 = self._make_layer(block, 64, layers[0])
#         self.layer2 = self._make_layer(block, 128, layers[1], stride=2)
#         self.layer3 = self._make_layer(block, 256, layers[2], stride=2)
#         self.layer4 = self._make_layer(
#             block, 512, layers[3], stride=last_stride)
#
#         self.attention1 = CBAM(64)
#         self.attention2 = CBAM(128)
#         self.attention3 = CBAM(256)
#         self.attention4 = CBAM(512)
#     def _make_layer(self, block, planes, blocks, stride=1):
#         downsample = None
#         if stride != 1 or self.inplanes != planes * block.expansion:
#             downsample = nn.Sequential(
#                 nn.Conv2d(self.inplanes, planes * block.expansion,
#                           kernel_size=1, stride=stride, bias=False),
#                 nn.BatchNorm2d(planes * block.expansion),
#             )
#
#         layers = []
#         layers.append(block(self.inplanes, planes, stride, downsample))
#         self.inplanes = planes * block.expansion
#         for i in range(1, blocks):
#             layers.append(block(self.inplanes, planes))
#
#         return nn.Sequential(*layers)
#
#     def forward(self, x):
#         x = self.conv1(x)
#         x = self.bn1(x)
#         # x = self.relu(x)    # add missed relu
#         x = self.maxpool(x)
#
#         x = self.layer1(x)
#         # x = self.attention1(x)
#         x = self.layer2(x)
#         x = self.attention2(x)
#         x = self.layer3(x)
#         x = self.attention3(x)
#         x = self.layer4(x)
#         x = self.attention4(x)
#
#         return x
#
#     def load_param(self, model_path):
#         param_dict = torch.load(model_path)
#         for i in param_dict:
#             if 'fc' in i:
#                 continue
#             self.state_dict()[i].copy_(param_dict[i])
#
#     def random_init(self):
#         for m in self.modules():
#             if isinstance(m, nn.Conv2d):
#                 n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
#                 m.weight.data.normal_(0, math.sqrt(2. / n))
#             elif isinstance(m, nn.BatchNorm2d):
#                 m.weight.data.fill_(1)
#                 m.bias.data.zero_()

from modeling.backbones.resnet_2 import BasicBlock2
class BasicConv2d(nn.Module):
    def __init__(self, in_planes, out_planes, kernel_size, stride, padding=0):
        super(BasicConv2d, self).__init__()
        self.conv = nn.Conv2d(in_planes, out_planes,
                              kernel_size=kernel_size, stride=stride,
                              padding=padding, bias=False) # verify bias false
        self.bn = nn.BatchNorm2d(out_planes)
        self.relu = nn.ReLU(inplace=True)

    def forward(self, x):
        x = self.conv(x)
        x = self.bn(x)
        x = self.relu(x)
        return x
def init_params(x):

    if x is None:
        return

    for m in x.modules():
        if isinstance(m, nn.Conv2d):
            nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
            if m.bias is not None:
                nn.init.constant_(m.bias, 0)
        elif isinstance(m, nn.BatchNorm2d):
            nn.init.normal_(m.weight, 1, 0.02)
            nn.init.constant_(m.bias, 0)
        elif isinstance(m, nn.BatchNorm1d):
            nn.init.normal_(m.weight, 1, 0.02)
            nn.init.constant_(m.bias, 0)
        elif isinstance(m, nn.Linear):
            # nn.init.normal_(m.weight, 0, 0.01)
            nn.init.kaiming_normal_(m.weight, a=0, mode='fan_out', nonlinearity='relu')
            if m.bias is not None:
                nn.init.constant_(m.bias, 0)
def weights_init_kaiming(m):
    classname = m.__class__.__name__
    if classname.find('Linear') != -1:
        nn.init.kaiming_normal_(m.weight, a=0, mode='fan_out')
        nn.init.constant_(m.bias, 0.0)
    elif classname.find('Conv') != -1:
        nn.init.kaiming_normal_(m.weight, a=0, mode='fan_in')
        if m.bias is not None:
            nn.init.constant_(m.bias, 0.0)
    elif classname.find('BatchNorm') != -1:
        if m.affine:
            nn.init.constant_(m.weight, 1.0)
            nn.init.constant_(m.bias, 0.0)
def weights_init_classifier(m):
    classname = m.__class__.__name__
    if classname.find('Linear') != -1:
        nn.init.normal_(m.weight, std=0.001)
        if m.bias:
            nn.init.constant_(m.bias, 0.0)
class MultiScaleLayer_v2(nn.Module):
    def __init__(self, in_planes, num_classes):
        super(MultiScaleLayer_v2, self).__init__()
        self.scale1 = BasicConv2d(in_planes//4, in_planes//4, (1,3), 1, padding=(0,1))
        self.scale2 = BasicConv2d(in_planes//4, in_planes//4, (3,1), 1, padding=(1,0))
        self.scale3 = BasicConv2d(in_planes//4, in_planes//4, (1,5), 1, padding=(0,2))
        self.scale4 = BasicConv2d(in_planes//4, in_planes//4, (5,1), 1, padding=(2,0))

        # from .osnet import OSBlock
        # self.block = OSBlock(in_planes, in_planes)
        # init_params(self.block)

        init_params(self.scale1)
        init_params(self.scale2)
        init_params(self.scale3)
        init_params(self.scale4)

        # from layers.gem_pool import GeneralizedMeanPoolingP
        # self.gap = GeneralizedMeanPoolingP()
        self.gap = nn.AdaptiveAvgPool2d(1)
        self.bottleneck = nn.BatchNorm1d(in_planes)
        self.bottleneck.bias.requires_grad_(False)
        self.classifier = nn.Linear(in_planes, num_classes, bias=False)
        self.bottleneck.apply(weights_init_kaiming)
        self.classifier.apply(weights_init_classifier)

    def forward(self, x):
        x1, x2, x3, x4 = torch.chunk(x, 4, dim=1)
        x1 = self.scale1(x1)
        x2 = self.scale2(x2)
        x3 = self.scale3(x3)
        x4 = self.scale4(x4)
        x = torch.cat((x1, x2, x3, x4), dim=1)
        x = self.gap(x)
        x = x.view(-1, x.size()[1])
        x = self.bottleneck(x)  # normalize for angular softmax
        x = self.classifier(x)
        return x
class ResNet_CBAM_group(nn.Module):
    def __init__(self, last_stride=2, layers=[3, 4, 6, 3]):
        self.inplanes = 64
        super().__init__()
        self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3,
                               bias=False)
        self.bn1 = nn.BatchNorm2d(64)
        # self.relu = nn.ReLU(inplace=True)   # add missed relu
        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
        self.layer1 = self._make_layer(BasicBlock, 64, layers[0])
        self.layer2 = self._make_layer(BasicBlock, 128, layers[1], stride=2)
        self.layer3 = self._make_layer(BasicBlock, 256, layers[2], stride=2)
        self.layer4 = self._make_layer(
            # BasicBlock,
            BasicBlock_group,
            512, layers[3], stride=last_stride)

        # self.attention1 = CA(64,64)
        # self.attention2 = CA(128,128)
        # self.attention3 = CA(256,256)
        # self.attention4 = CA(512,512)

        # self.attention1 = CBAM(64)
        # self.attention2 = CBAM(128)


        # self.attention3 = CBAM(256)
        # self.attention3 = CBAM_fby(256)
        # self.attention3 = CBAM_fby_cat(256)


        # self.attention4 = CBAM(512)

        # self.attention1 = CBAM_fby(64)
        # self.attention2 = CBAM_fby(128)
        # self.attention3 = CBAM_fby(256)
        # self.attention4 = CBAM_fby(512)

        self.attention1 = GlobalContextBlock_fby(inplanes=64, ratio=0.25)
        # self.attention2 = GlobalContextBlock(inplanes=128, ratio=0.25)
        # self.attention3 = GlobalContextBlock(inplanes=256, ratio=0.25)
        # self.attention4 = GlobalContextBlock(inplanes=512, ratio=0.25)
    def _make_layer(self, block, planes, blocks, stride=1):
        downsample = None
        if stride != 1 or self.inplanes != planes * block.expansion:
            downsample = nn.Sequential(
                nn.Conv2d(self.inplanes, planes * block.expansion,
                          kernel_size=1, stride=stride, bias=False),
                nn.BatchNorm2d(planes * block.expansion),
            )

        layers = []
        layers.append(block(self.inplanes, planes, stride, downsample))
        self.inplanes = planes * block.expansion
        for i in range(1, blocks):
            layers.append(block(self.inplanes, planes))

        return nn.Sequential(*layers)

    def forward(self, x):
        x = self.conv1(x)
        x = self.bn1(x)
        # x = self.relu(x)    # add missed relu
        x = self.maxpool(x)

        x = self.layer1(x)
        x = self.attention1(x)
        x = self.layer2(x)
        # x = self.attention2(x)
        x = self.layer3(x)
        # x = self.attention3(x)
        x = self.layer4(x)
        # x = self.attention4(x)

        return x

    def load_param(self, model_path):
        self.random_init()
        param_dict = torch.load(model_path)
        for i in param_dict:
            if 'fc' in i:
                continue
            aa = param_dict[i]
            bb = self.state_dict()[i]
            if not aa.shape == bb.shape:
                continue
            self.state_dict()[i].copy_(param_dict[i])

    def random_init(self):
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
                m.weight.data.normal_(0, math.sqrt(2. / n))
            elif isinstance(m, nn.BatchNorm2d):
                m.weight.data.fill_(1)
                m.bias.data.zero_()
if __name__ == "__main__":
    import time
    # mode = ResNet_CBAM_group().cuda()
    # input= torch.randn(1,3,248,124).cuda()

    mode = ResNet_CBAM_group()
    input= torch.randn(1,3,248,124)

    start = time.time()
    mode(input)
    end = time.time()
    print((end - start)*1000)
